# _*_ coding:utf-8 _*_
# @Time      : 8:54
# @Author    :baizhou
from sklearn import svm
from . import ConverImage
import numpy as np
import os
import cv2
from sklearn.externals import joblib

from sklearn.multiclass import OneVsRestClassifier


def resize_image(image, height=128, width=128):
    top, bottom, left, right = (0, 0, 0, 0)
    # 获取图像尺寸
    h, w, _ = image.shape
    # 对于长宽不相等的图片，找到最长的一边
    longest_edge = max(h, w)
    # 计算短边需要增加多上像素宽度使其与长边等长
    if h < longest_edge:
        dh = longest_edge - h
        top = dh // 2
        bottom = dh - top
    elif w < longest_edge:
        dw = longest_edge - w
        left = dw // 2
        right = dw - left
    else:
        pass

    # RGB颜色
    BLACK = [0, 0, 0]
    # 给图像增加边界，是图片长、宽等长，cv2.BORDER_CONSTANT指定边界颜色由value指定
    constant = cv2.copyMakeBorder(image, top, bottom, left, right, cv2.BORDER_CONSTANT, value=BLACK)
    # 调整图片大小并返回
    return cv2.resize(constant, (height, width))


def svmTrainMain(trainPath):
    # 模型本地化的路径
    savePath = r'E:\freash\1\svmModel.h5'
    trainSet, trainLabel = ConverImage.load_dataset(trainPath)
    # testSet, testLable = load_dataset(testPath)
    # print(testLable)
    model = OneVsRestClassifier(svm.SVC(kernel='rbf'))
    clf = model.fit(trainSet, trainLabel)
    # 训练完成的模型本地化
    joblib.dump(clf, savePath)
    # pret = clf.predict(testSet)
    # trainAcc = clf.score(trainSet, trainLabel)
    # testAcc = clf.score(testSet, testLable)

    data = {
        'savePath': savePath
    }
    return data


def svmPrect(savePath, inputImagePath):
    '''
    对用户输入的图片进行于测
    :param savePath: 训练之后的模型路径
    :param inputImage: 待预测的图片地址
    :return: 返回预测结果
    '''
    # 加载本地训练好的模型
    images = []
    res = ''
    clf = joblib.load(savePath)
    image = cv2.imread(inputImagePath)
    image = resize_image(image, 128, 128)
    images.append(np.array(image).flatten())
    preResult = clf.predict(images)
    # 只进行单张图片测试
    if preResult[0] ==0:
        res += '二等'
    if preResult[0] == 1:
        res += '一等'
    if preResult[0] ==2:
        res += '特等'
    data = {
        'res':res
    }
    return data



if __name__ == '__main__':
    trainPath = r'E:\work\assay\dataset\dataset5(new)'
    testPath = r'E:\freash\1\test'
    data = svmTrainMain(trainPath=trainPath)
    inputImagePath = r'E:\freash\1\train\lower\DSC05720.JPG'
    res = svmPrect(data['savePath'],inputImagePath=inputImagePath)



